package com.zhang.data.mongo.core;

import com.mongodb.MongoClientSettings;
import com.mongodb.client.ClientSession;
import com.zhang.data.mongo.annotation.Collection;
import com.mongodb.client.AggregateIterable;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.result.InsertOneResult;
import org.bson.Document;
import org.bson.codecs.configuration.CodecProvider;
import org.bson.codecs.configuration.CodecRegistry;
import org.bson.codecs.pojo.PojoCodecProvider;
import org.bson.conversions.Bson;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.function.Consumer;

import static org.bson.codecs.configuration.CodecRegistries.*;

public class MongoUtil {

    private MongoClient client;

    private String database;

    private final ThreadLocal<ClientSession> threadSession = new ThreadLocal<>();


    public MongoUtil(MongoClient client, String database) {
        this.client = client;
        this.database = database;
    }

    public <T> InsertOneResult insetOne(T document, Class<T> documentClass){
        MongoCollection<T> collection=this.getCollection(documentClass);
        ClientSession session = threadSession.get();
        if(null == session){
            return collection.insertOne(document);
        }
        InsertOneResult result= collection.insertOne(threadSession.get(),document);
        return result;
    }


    public <T> List<T> find(Bson document, Class<T> documentClass){
        MongoCollection<T> collection=this.getCollection(documentClass);
        System.out.println(collection.getCodecRegistry());
        List<T> list = new LinkedList<>();
        collection.find(document,documentClass).forEach(t -> list.add(t));
        return list;
    }

    public <T> List<T> find(Bson document, String collectionName, ResultMapper<T> mapper){
        MongoCollection<Document> collection=client.getDatabase(database).getCollection(collectionName);
        List<T> list = new LinkedList<>();
        collection.find(document).forEach(t -> list.add(mapper.map(t)));
        return list;
    }

    public <T> void findConsumer(Bson document, Class<T> documentClass, Consumer<T> consumer){
        MongoCollection<T> collection=this.getCollection(documentClass);
        collection.find(document,documentClass).forEach(consumer);
    }

    public <T> List<T> agg(List<? extends Bson> pipeline,String collectionName, Class<T> documentClass){
        MongoCollection<Document> collection=client.getDatabase(database).getCollection(collectionName);
        AggregateIterable<T> result= collection.aggregate(pipeline,documentClass);
        List<T> list = new LinkedList<>();
        result.forEach(t -> list.add(t));
        return list;
    }

    public Document agg(List<? extends Bson> pipeline,String collectionName){
        MongoCollection<Document> collection=client.getDatabase(database).getCollection(collectionName);
        AggregateIterable<Document> result= collection.aggregate(pipeline,Document.class);
        Document document = new Document();
        List<Document> list = new ArrayList<>();
        result.forEach(t -> list.add(t));
        document.append("result",list);
        return document;
    }

    public void drop(String collectionName){
        MongoCollection<Document> collection=client.getDatabase(database).getCollection(collectionName);
        collection.drop();
    }

    public <T> void drop(Class<T> documentClass){
        MongoCollection<T> collection=this.getCollection(documentClass);
        collection.drop();
    }

    private <T>  MongoCollection<T> getCollection(Class<T> documentClass){
        Collection collection1= documentClass.getAnnotation(Collection.class);
        String collectionName= documentClass.getSimpleName().toLowerCase();
        if (null != collection1){
            collectionName =collection1.name();
        }
        MongoCollection<T> collection=client.getDatabase(database)
                .getCollection(collectionName,documentClass);
        ClientSession clientSession=client.startSession();
        clientSession.startTransaction();
        return  collection;
    }

    private CodecRegistry getCodecRegistry(Class<?> documentClass){
        CodecProvider pojoCodecProvider = PojoCodecProvider.builder().register(documentClass).build();
        CodecRegistry pojoCodecRegistry = fromRegistries(
                MongoClientSettings.getDefaultCodecRegistry(),
                fromProviders(pojoCodecProvider));
        return pojoCodecRegistry;
    }

    public void startTransaction(){
        ClientSession clientSession=client.startSession();
        clientSession.startTransaction();
        threadSession.set(clientSession);
    }

    public void commitTransaction(){
       this.endTransaction(true);
    }


    public void abortTransaction(){
        this.endTransaction(false);
    }

    private void endTransaction(boolean isCommit){
        ClientSession clientSession=threadSession.get();
        if (null == clientSession){
            throw new IllegalStateException("There is no session exits");
        }
        if (isCommit){
            clientSession.commitTransaction();
        }else {
            clientSession.abortTransaction();
        }
        clientSession.close();
        threadSession.remove();
    }

}
